from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from starlette.websockets import WebSocketDisconnect
import uvicorn
import cv2
import base64
import json
from typing import Dict, Optional
import asyncio
from pydantic import BaseModel

from .camera import CameraManager
# from .ai_detector import AIDetector  # 暂时注释掉
from config.settings import *

def cleanup():
    try:
        if hasattr(app.state, 'camera_manager'):
            for camera in app.state.camera_manager.get_all_cameras().values():
                if camera:
                    camera.stop()
    except Exception as e:
        print(f"清理摄像头资源时出错: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时执行
    camera_manager = CameraManager()
    camera_manager.init_cameras()
    app.state.camera_manager = camera_manager
    yield
    # 关闭时执行
    cleanup()

app = FastAPI(lifespan=lifespan)

# 挂载静态文件
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

class CameraAction(BaseModel):
    action: str

@app.get("/")
async def get_index():
    with open(BASE_DIR / "frontend" / "index.html", encoding='utf-8') as f:
        return HTMLResponse(f.read())

@app.get("/api/cameras")
async def get_cameras():
    """获取所有摄像头状态"""
    cameras = app.state.camera_manager.get_all_cameras()
    return {"cameras": [{"id": id, "status": "active" if camera else "inactive"} 
            for id, camera in cameras.items()]}

@app.post("/api/cameras/{camera_id}/control")
async def control_camera(camera_id: int, action: CameraAction):
    if action.action == "start":
        if app.state.camera_manager.add_camera(camera_id):
            camera = app.state.camera_manager.get_camera(camera_id)
            if camera and camera.is_opened():
                return {"status": "success"}
            app.state.camera_manager.remove_camera(camera_id)
            raise HTTPException(status_code=400, detail=f"摄像头 {camera_id} 无法访问或已被占用")
        raise HTTPException(status_code=400, detail=f"摄像头 {camera_id} 添加失败")
    elif action.action == "stop":
        if app.state.camera_manager.remove_camera(camera_id):
            return {"status": "success"}
        raise HTTPException(status_code=400, detail=f"摄像头 {camera_id} 不存在")
    raise HTTPException(status_code=400, detail="无效的操作指令")

@app.websocket("/ws/video/{camera_id}")
async def video_websocket(websocket: WebSocket, camera_id: int):
    await websocket.accept()
    
    try:
        camera = app.state.camera_manager.get_camera(camera_id)
        if not camera:
            await websocket.close(code=1000, reason="Camera not found")
            return

        while True:
            try:
                frame = camera.get_frame()
                if frame is not None:
                    _, buffer = cv2.imencode('.jpg', frame)
                    data = {
                        "image": base64.b64encode(buffer).decode('utf-8'),
                    }
                    await websocket.send_text(json.dumps(data))
                await asyncio.sleep(1/FRAME_RATE)
            except WebSocketDisconnect:
                print(f"客户端断开连接: camera_id={camera_id}")
                break
            except asyncio.CancelledError:
                print(f"WebSocket任务被取消: camera_id={camera_id}")
                break
            except Exception as e:
                print(f"处理帧时出错: {e}")
                break
    except Exception as e:
        print(f"WebSocket错误: {e}")
    finally:
        print(f"清理摄像头资源: camera_id={camera_id}")
        try:
            await websocket.close()
        except:
            pass

@app.websocket("/ws/camera/{camera_id}")
async def websocket_endpoint(websocket: WebSocket, camera_id: int):
    try:
        await websocket.accept()
        camera = app.state.camera_manager.get_camera(camera_id)
        if not camera:
            await websocket.close(code=1000, reason="Camera not found")
            return
            
        while True:
            frame = camera.get_frame()
            if frame is not None:
                _, buffer = cv2.imencode('.jpg', frame)
                await websocket.send_bytes(buffer.tobytes())
            await asyncio.sleep(0.033)
            
    except WebSocketDisconnect:
        pass
    except Exception as e:
        pass
    finally:
        try:
            await websocket.close()
        except:
            pass

if __name__ == "__main__":
    try:
        config = uvicorn.Config("backend.main:app", host=HOST, port=PORT, reload=True)
        server = uvicorn.Server(config)
        
        # 注册清理函数
        import atexit
        atexit.register(cleanup)
        
        server.run()
    except KeyboardInterrupt:
        cleanup()